#This module represents the actual field. It solves the 1D Richards equation for a 0.32m sandy clay
#loam field using the first input provided by the open-loop mpc solution

import numpy as np 
from parameters_mpc import temporalVariable, spatialVariables_1D, Sandyclayloam
from richards_equation_1D import Richards1D_vectorized
from scipy import integrate

#Load the  soil parameters, spatial, and the temporal variables
soilPars = Sandyclayloam()
totalDepth, axialNodes, axialDistance, totalNodes, numOfActuators, interPoints= spatialVariables_1D()  #Spatial parameters
samplingTime, samplingTimeInternal, internalTimeSteps, totTimeSteps = temporalVariable() #time parameters


def ODE(head, timeArray, irrigAmount, cropCoeff, refEvap, soilPars):
    return Richards1D_vectorized(head, irrigAmount, cropCoeff, refEvap, soilPars)


def SimulateTheActualPlant(initialCondition, prescribedIrrigation, referenceEvap, cropCoeff):
    
    irrigationSequence = np.zeros(totTimeSteps)
    irrigationSequence[0] = prescribedIrrigation
    
    cropCoeffSequence = cropCoeff*np.ones(totTimeSteps)
    refEvapSequence = referenceEvap*np.ones(totTimeSteps)
    
    headArray = np.zeros((totTimeSteps+1, totalNodes))
    headArray[0,:] = initialCondition
    
    i=1
    sol = integrate.solve_ivp(fun=lambda t, y: ODE(y, t, irrigationSequence[i-1],cropCoeffSequence[i-1],refEvapSequence[i-1], soilPars),
                              t_span=[0,samplingTimeInternal],y0=tuple(headArray[i-1]),method='LSODA')
    
    headArray[i,:]=sol.y[:,-1]
    
    return headArray[-1,:]
    

    